modelopt/torch/speculative/eagle/conversion.py (1)

35-39: Avoid relying on private registry internals; register with a stable key and fail fast when unmapped.

• Using EagleDMRegistry._registry directly is a private detail and brittle.
• The key "base_model_class" is opaque; prefer a stable, informative key (e.g., fully-qualified class name).
• If no base class match is found, EagleDMRegistry.convert will raise a KeyError later with less context. Raise a clear error when the mapping cannot be inferred.

Apply this diff to make the registration more robust and explicit:

-    if original_cls not in EagleDMRegistry:
-        for cls in EagleDMRegistry._registry:
-            if issubclass(original_cls, cls):
-                EagleDMRegistry.register({original_cls: "base_model_class"})(EagleDMRegistry[cls])
-                break
+    if original_cls not in EagleDMRegistry:
+        # Register a mapping for subclasses to the same dynamic class as their base.
+        found = False
+        for base_cls in list(EagleDMRegistry._registry.keys()):
+            if issubclass(original_cls, base_cls):
+                key = f"{original_cls.__module__}.{original_cls.__name__}"
+                EagleDMRegistry.register({original_cls: key})(EagleDMRegistry[base_cls])
+                found = True
+                break
+        if not found:
+            raise KeyError(
+                f"No Eagle dynamic mapping for {original_cls.__module__}.{original_cls.__name__}. "
+                f"Ensure a compatible base class is registered in EagleDMRegistry."
+            )

modelopt/torch/speculative/plugins/megatron_eagle.py (7)

748-751: Disable sequence_parallel in offline mode — add a guard log.

For clarity during debugging, emit a one-time info when flipping sequence_parallel to False in offline runs so users aren’t surprised by the override.

Apply this diff:

         # sequence_parallel is not used in offline eagle
         if self.eagle_offline:
-            self.config.sequence_parallel = False
+            if self.config.sequence_parallel:
+                warnings.warn("EAGLE offline: forcing sequence_parallel = False")
+            self.config.sequence_parallel = False

---

768-768: Update comment to reflect “offline” (detached class no longer exists).

Replace “detached eagle” with “offline eagle” to avoid confusion.

-        # layer ids are not used in detached eagle, but we need to set this to have correct fc_input_size_multiplier
+        # Layer IDs are not used in offline eagle, but we set them to get the correct fc_input_size_multiplier

---

853-866: Offline path: assert input shape assumptions and document expectations.

In offline mode, _get_eagle_input_hidden_states expects aux_hidden_states to already be concatenated ([s, b, k*h]). Add explicit asserts to catch silent shape mismatches.

     def _get_eagle_input_hidden_states(self, hidden_states: torch.Tensor, apply_fc: bool = True):
         """When _aux_hidden_states is not empty, then this is EAGLE-3.
@@
-        if not self.eagle_offline:
+        if not self.eagle_offline:
             if len(self._aux_hidden_states) == 0:
                 return hidden_states
@@
-        if apply_fc:
+        if apply_fc:
+            # In offline mode, hidden_states may already be [s, b, k*h]; validate k.
+            if self.eagle_offline:
+                h = self.config.hidden_size
+                assert hidden_states.shape[-1] % h == 0, (
+                    f"Expected aux_hidden_states hidden dim to be a multiple of {h}, "
+                    f"got {hidden_states.shape[-1]}"
+                )
             # [s / TP, b, 3h] -> [s / TP, b, h]
             return self.eagle_module.fc(hidden_states)[0]

---

1076-1086: Make loss alignment explicit and robust across online/offline lengths.

Inferring offline via labels.shape[1] < eagle_logits.shape[0] is brittle. Align by checking the delta and fail fast on unexpected shapes.

         if self.eagle_self_logit_distillation:
             mapping = self.eagle_module.d2t if hasattr(self.eagle_module, "d2t") else None
             token_loss = self.kld(eagle_logits[:-1, :, :], logits[1:, :, :], mapping)
-        elif labels.shape[1] < eagle_logits.shape[0]:
-            token_loss = self.compute_language_model_loss(labels[:, 1:], eagle_logits[:-2, :, :])
-        else:
-            token_loss = self.compute_language_model_loss(labels[:, 1:], eagle_logits[:-1, :, :])
+        else:
+            target_len = labels.shape[1] - 1
+            pred_len = eagle_logits.shape[0]
+            if pred_len - 1 == target_len:
+                aligned = eagle_logits[:-1, :, :]
+            elif pred_len - 2 == target_len:
+                aligned = eagle_logits[:-2, :, :]
+            else:
+                raise ValueError(
+                    f"Unexpected lengths for EAGLE loss: labels={labels.shape}, "
+                    f"eagle_logits={eagle_logits.shape}"
+                )
+            token_loss = self.compute_language_model_loss(labels[:, 1:], aligned)

---

1298-1301: Consolidate accuracy-slice logic to avoid off‑by‑one drift.

Four places differ only by the starting offset and offline end trim (-1 vs -2). Centralize this to a small helper to prevent future inconsistencies.

Add this helper (outside forward):

def _slice_for_acc(self, logits_1d: torch.Tensor, start: int) -> torch.Tensor:
    end = -2 if self.eagle_offline else -1
    return logits_1d[start:end, :, :]

Then replace:

• eagle_logits_0[:-2 or :-1] -> self._slice_for_acc(eagle_logits_0, 0)
• eagle_logits_1[1:-2 or 1:-1] -> self._slice_for_acc(eagle_logits_1, 1)
• eagle_logits_2[2:-2 or 2:-1] -> self._slice_for_acc(eagle_logits_2, 2)
• eagle_logits_3[3:-2 or 3:-1] -> self._slice_for_acc(eagle_logits_3, 3)

Also applies to: 1337-1343, 1379-1385, 1421-1426

---

1245-1249: Avoid eager .cpu() copies on hot paths.

If return_eagle_inputs is used frequently, consider non_blocking transfers to CPU to minimize stalls.

-            return {
-                "input_ids": input_ids.squeeze(0).cpu(),
-                "aux_hidden_states": eagle_module_input_hidden_states.squeeze(1).cpu(),
-                "hidden_states": hidden_states.squeeze(1).cpu(),
-            }
+            return {
+                "input_ids": input_ids.squeeze(0).to("cpu", non_blocking=True),
+                "aux_hidden_states": eagle_module_input_hidden_states.squeeze(1).to(
+                    "cpu", non_blocking=True
+                ),
+                "hidden_states": hidden_states.squeeze(1).to("cpu", non_blocking=True),
+            }

---

1461-1473: Minor: clear hooks along early returns.

You clear self._aux_hidden_states on early-return with no children. If forward hooks were registered conditionally, also consider guarding to avoid accumulating state across calls.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

modelopt/torch/speculative/plugins/megatron_eagle.py (1)

1196-1211: Addressed prior review: offline forward now validates inputs and forbids return_eagle_inputs.

This resolves the earlier “NoneType later” hazard and ambiguous semantics in offline mode.

modelopt/torch/speculative/plugins/megatron_eagle.py (4)

748-751: Disable sequence_parallel in offline: good; add an explicit user-facing warning when flipping it.

Silent mutation of config can surprise callers. Emit a warning if it was True.

Apply:

         # sequence_parallel is not used in offline eagle
         if self.eagle_offline:
-            self.config.sequence_parallel = False
+            if self.config.sequence_parallel:
+                warnings.warn("EAGLE offline mode: forcibly disabling sequence_parallel.", stacklevel=1)
+            self.config.sequence_parallel = False

---

767-768: Update terminology: “detached” → “offline”.

Comment still mentions “detached eagle”; recommend aligning to “offline” to avoid confusion.

-        # layer ids are not used in detached eagle, but we need to set this to have correct fc_input_size_multiplier
+        # layer ids are not used in offline EAGLE, but we need to set this to have correct fc_input_size_multiplier

---

853-866: Offline/path gate in _get_eagle_input_hidden_states: OK. Consider defensive clearing.

Looks correct. As a safety edge-case when toggling modes within a long-lived object, consider clearing _aux_hidden_states on offline entry to avoid accidental reuse.

     def _get_eagle_input_hidden_states(self, hidden_states: torch.Tensor, apply_fc: bool = True):
@@
-        if not self.eagle_offline:
+        if not self.eagle_offline:
             if len(self._aux_hidden_states) == 0:
                 return hidden_states
             # [s / TP, b, len(self._aux_hidden_states) * h]
             hidden_states = torch.cat(self._aux_hidden_states, dim=-1)
             self._aux_hidden_states.clear()
+        else:
+            # Ensure no stale states leak in if mode was toggled.
+            self._aux_hidden_states.clear()

---

1083-1086: Tighten condition for offline label/logit alignment.

Use equality instead of less-than to avoid mis-slicing when labels are unexpectedly shorter.

-        elif labels.shape[1] < eagle_logits.shape[0]:
+        elif labels.shape[1] + 1 == eagle_logits.shape[0]:
             token_loss = self.compute_language_model_loss(labels[:, 1:], eagle_logits[:-2, :, :])

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

modelopt/torch/speculative/plugins/transformers.py (1)

692-701: Use inference_mode for read-only blocks.

Replace torch.no_grad() with torch.inference_mode() where no tensors require grad to reduce dispatcher overhead and enable viewless inference.

Example:

-        with torch.no_grad() if freeze_base_model else contextlib.nullcontext():
+        with torch.inference_mode() if freeze_base_model else contextlib.nullcontext():

And similarly for the later embed-only sections.

Also applies to: 842-853, 874-888, 922-941, 965-984

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

modelopt/torch/speculative/plugins/megatron_eagle.py (1)

1199-1211: Fix offline inputs: allow EAGLE‑1, validate args, and centralize PP early‑return.

Offline forward still requires aux_hidden_states and places the PP early‑return only in the online branch. This blocks EAGLE‑1 offline and can NPE when aux is absent. Move the PP early‑return out of the if/else and only require aux when use_aux_hidden_state=True. Also select src_states accordingly.

Apply this diff:

@@
-        if self.eagle_offline:
-            # aux_hidden_states and hidden_states are provided for offline eagle
-            # _base_model_forward is skipped
-            if return_eagle_inputs:
-                raise ValueError("return_eagle_inputs is unsupported in EAGLE offline mode.")
-            aux_hidden_states = kwargs.get("aux_hidden_states")
-            hidden_states = kwargs.get("hidden_states")
-            if aux_hidden_states is None or hidden_states is None:
-                raise ValueError(
-                    "EAGLE offline mode requires kwargs: aux_hidden_states=[s,b,k*h], "
-                    "hidden_states=[s,b,h]."
-                )
+        if self.eagle_offline:
+            # hidden_states (required) and aux_hidden_states (optional for EAGLE‑1) are provided for offline EAGLE.
+            # _base_model_forward is skipped
+            if return_eagle_inputs:
+                raise ValueError("return_eagle_inputs is unsupported in EAGLE offline mode.")
+            hidden_states = kwargs.get("hidden_states", None)
+            if hidden_states is None:
+                raise ValueError("EAGLE offline mode requires kwargs: hidden_states=[s,b,h].")
+            aux_hidden_states = kwargs.get("aux_hidden_states", None)
+            if self.eagle_config.use_aux_hidden_state and aux_hidden_states is None:
+                raise ValueError(
+                    "EAGLE‑3 offline additionally requires aux_hidden_states=[s,b,k*h] when "
+                    "use_aux_hidden_state=True."
+                )
         else:
@@
-            # Typically, this is only the case when PP > 1.
-            if not self.post_process:
-                return hidden_states
+        # Typically, this is only the case when PP > 1.
+        if not self.post_process:
+            return hidden_states
@@
-        if self.eagle_offline:
-            eagle_module_input_hidden_states = self._get_eagle_input_hidden_states(
-                aux_hidden_states, apply_fc=self.eagle_config.use_aux_hidden_state
-            )
+        if self.eagle_offline:
+            src_states = (
+                aux_hidden_states if self.eagle_config.use_aux_hidden_state else hidden_states
+            )
+            eagle_module_input_hidden_states = self._get_eagle_input_hidden_states(
+                src_states, apply_fc=self.eagle_config.use_aux_hidden_state
+            )

Also applies to: 1226-1229, 1235-1238

modelopt/torch/speculative/plugins/megatron_eagle.py (6)

1241-1256: Preserve batch dims for return_eagle_inputs (avoid squeeze assumptions).

Squeezing dim 1 breaks when batch_size > 1. Return tensors with their batch dims intact.

-            return {
-                "input_ids": input_ids.squeeze(0).cpu(),
-                "aux_hidden_states": eagle_module_input_hidden_states.squeeze(1).cpu(),
-                "hidden_states": hidden_states.squeeze(1).cpu(),
-            }
+            return {
+                "input_ids": input_ids.cpu(),
+                "aux_hidden_states": eagle_module_input_hidden_states.cpu(),
+                "hidden_states": hidden_states.cpu(),
+            }

---

767-773: Update comment: “detached eagle” wording is outdated.

This PR removes DetachedEagle; update the comment to reflect “offline EAGLE” semantics.

-        # layer ids are not used in detached eagle, but we need to set this to have correct fc_input_size_multiplier
+        # Layer IDs are not used in offline EAGLE‑1, but we set them for correct fc_input_size_multiplier in EAGLE‑3.

---

1083-1085: Make offline loss branch explicit (avoid shape‑heuristic).

Using labels.shape[1] < eagle_logits.shape[0] is brittle. Gate by eagle_offline and optionally assert expected alignment.

-        elif labels.shape[1] < eagle_logits.shape[0]:
+        elif self.eagle_offline:
+            assert eagle_logits.shape[0] >= labels.shape[1] + 1, \
+                "Offline EAGLE expects logits length >= labels length + 1."
             token_loss = self.compute_language_model_loss(labels[:, 1:], eagle_logits[:-2, :, :])

---

1309-1311: DRY the offline/online top‑1 slicing logic.

The repeated ternaries ([:-2] vs. [:-1]) are easy to drift. Factor a small helper to slice logits consistently.

+    def _slice_for_top1(self, logits):
+        return logits[:-2, :, :] if self.eagle_offline else logits[:-1, :, :]
@@
-                gathered_logits = gather_from_tensor_model_parallel_region(
-                    eagle_logits_0[:-2, :, :] if self.eagle_offline else eagle_logits_0[:-1, :, :]
-                )
+                gathered_logits = gather_from_tensor_model_parallel_region(
+                    self._slice_for_top1(eagle_logits_0)
+                )
@@
-                gathered_logits = gather_from_tensor_model_parallel_region(
-                    eagle_logits_1[1:-2, :, :] if self.eagle_offline else eagle_logits_1[1:-1, :, :]
-                )
+                gathered_logits = gather_from_tensor_model_parallel_region(
+                    self._slice_for_top1(eagle_logits_1[1:, :, :])
+                )
@@
-                gathered_logits = gather_from_tensor_model_parallel_region(
-                    eagle_logits_2[2:-2, :, :] if self.eagle_offline else eagle_logits_2[2:-1, :, :]
-                )
+                gathered_logits = gather_from_tensor_model_parallel_region(
+                    self._slice_for_top1(eagle_logits_2[2:, :, :])
+                )
@@
-                gathered_logits = gather_from_tensor_model_parallel_region(
-                    eagle_logits_3[3:-2, :, :] if self.eagle_offline else eagle_logits_3[3:-1, :, :]
-                )
+                gathered_logits = gather_from_tensor_model_parallel_region(
+                    self._slice_for_top1(eagle_logits_3[3:, :, :])
+                )

Also applies to: 1350-1351, 1392-1393, 1434-1435

---

846-866: Online/Offline hidden state selection looks correct; small doc tweak.

Logic to bypass aux accumulation in offline is right. Consider noting in the docstring that offline passes either aux (EAGLE‑3) or base hidden_states (EAGLE‑1).

---

1251-1251: Optional: avoid unnecessary full TP gather of base logits on return_eagle_inputs.

If the only consumer is offline precompute, gathering logits_sbh here is unnecessary. Returning CPU’d inputs without TP gather can reduce overhead.

If you keep it, please confirm a downstream consumer requires full‑vocab logits. Otherwise, drop this gather.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

modelopt/torch/speculative/plugins/megatron_eagle.py (3)

1199-1211: Offline EAGLE-1 is blocked; make aux_hidden_states optional and validate by config.

Requiring aux_hidden_states unconditionally prevents offline EAGLE‑1. Gate the requirement on use_aux_hidden_state and always require hidden_states. Also reject return_eagle_inputs early (kept).

Apply:

         if self.eagle_offline:
             # aux_hidden_states and hidden_states are provided for offline eagle
             # _base_model_forward is skipped
             if return_eagle_inputs:
                 raise ValueError("return_eagle_inputs is unsupported in EAGLE offline mode.")
-            aux_hidden_states = kwargs.get("aux_hidden_states")
-            hidden_states = kwargs.get("hidden_states")
-            if aux_hidden_states is None or hidden_states is None:
-                raise ValueError(
-                    "EAGLE offline mode requires kwargs: aux_hidden_states=[s,b,k*h], "
-                    "hidden_states=[s,b,h]."
-                )
+            hidden_states = kwargs.get("hidden_states", None)
+            if hidden_states is None:
+                raise ValueError("EAGLE offline mode requires kwargs: hidden_states=[s,b,h].")
+            aux_hidden_states = kwargs.get("aux_hidden_states", None)
+            if self.eagle_config.use_aux_hidden_state and aux_hidden_states is None:
+                raise ValueError(
+                    "EAGLE‑3 offline additionally requires aux_hidden_states=[s,b,k*h] when "
+                    "use_aux_hidden_state=True."
+                )

---

1226-1228: Centralize PP early‑return so it runs for both online and offline paths.

Right now the early return is only reachable in the online branch. Move it out so offline PP does not proceed to output/loss.

-            # Typically, this is only the case when PP > 1.
-            if not self.post_process:
-                return hidden_states
+        # Typically, this is only the case when PP > 1.
+        if not self.post_process:
+            return hidden_states

---

1235-1238: Offline EAGLE input should come from aux or base hidden by config.

Use aux when use_aux_hidden_state=True (EAGLE‑3), otherwise pass hidden_states (EAGLE‑1). Also set apply_fc accordingly.

-        if self.eagle_offline:
-            eagle_module_input_hidden_states = self._get_eagle_input_hidden_states(
-                aux_hidden_states, apply_fc=self.eagle_config.use_aux_hidden_state
-            )
+        if self.eagle_offline:
+            src_states = (
+                aux_hidden_states
+                if self.eagle_config.use_aux_hidden_state
+                else hidden_states
+            )
+            eagle_module_input_hidden_states = self._get_eagle_input_hidden_states(
+                src_states, apply_fc=self.eagle_config.use_aux_hidden_state
+            )

modelopt/torch/speculative/plugins/megatron_eagle.py (4)

1306-1321: Guard accuracy reporting against empty slices (short sequences).

When seq_len is too short, slices like[:-2] can be empty → argmax and division by zero crash. Skip metrics if there is no token to score.

Apply to each block (1st→4th) similarly; example shown for the 1st block:

-            with torch.no_grad():
-                gathered_logits = gather_from_tensor_model_parallel_region(
-                    eagle_logits_0[:-2, :, :] if self.eagle_offline else eagle_logits_0[:-1, :, :]
-                )
-                eagle_top1 = gathered_logits.transpose(0, 1).argmax(dim=-1)
-                if self.eagle_config.draft_vocab_size != self.eagle_config.vocab_size:
-                    eagle_top1 += self.eagle_module.d2t[eagle_top1]
-                top1_p = torch.eq(labels[:, 1:], eagle_top1).sum() / eagle_top1.numel()
-                acc.append(top1_p)
+            with torch.no_grad():
+                sl = eagle_logits_0.shape[0] - (2 if self.eagle_offline else 1)
+                if sl > 0:
+                    logits_slice = (
+                        eagle_logits_0[:sl, :, :]
+                    )
+                    gathered_logits = gather_from_tensor_model_parallel_region(logits_slice)
+                    eagle_top1 = gathered_logits.transpose(0, 1).argmax(dim=-1)
+                    if self.eagle_config.draft_vocab_size != self.eagle_config.vocab_size:
+                        eagle_top1 += self.eagle_module.d2t[eagle_top1]
+                    denom = max(1, eagle_top1.numel())
+                    top1_p = torch.eq(labels[:, 1:], eagle_top1).sum() / denom
+                    acc.append(top1_p)

Repeat the same shape guard for the 2nd, 3rd, and 4th accuracy blocks.

Also applies to: 1348-1363, 1390-1405, 1432-1447

---

1076-1086: Loss shape handling: prefer explicit offline flag over shape heuristics.

Using labels.shape[1] < eagle_logits.shape[0] to infer offline mode is brittle. Gate the[:-2] slice with self.eagle_offline for clarity and safety.

-        elif labels.shape[1] < eagle_logits.shape[0]:
-            token_loss = self.compute_language_model_loss(labels[:, 1:], eagle_logits[:-2, :, :])
-        else:
-            token_loss = self.compute_language_model_loss(labels[:, 1:], eagle_logits[:-1, :, :])
+        else:
+            slice_ = eagle_logits[:-2, :, :] if self.eagle_offline else eagle_logits[:-1, :, :]
+            token_loss = self.compute_language_model_loss(labels[:, 1:], slice_)

---

768-771: Stale terminology: “detached eagle” in comment.

Detached wrapper was removed; update the comment to “offline EAGLE”.

-        # layer ids are not used in detached eagle, but we need to set this to have correct fc_input_size_multiplier
+        # Layer ids are not used in offline EAGLE, but we set them to derive the correct fc_input_size_multiplier.

---

640-647: Comment contradicts code: clarify EAGLE‑1/‑3 inputs.

Code concatenates [embeddings, hidden_states] → 2h then fc for EAGLE‑1; EAGLE‑3 uses aux + specialized attention to reduce 2h→h. Update the comment.

-        # EAGLE-1 uses [s, b, h] input but EAGLE-3 uses [s, b, 2h] input
+        # EAGLE‑1: concat [embeddings, hidden_states] → [s,b,2h], then fc → [s,b,h].
+        # EAGLE‑3: consumes aux/hidden; first attention reduces 2h → h (specialized qkv).

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

modelopt/torch/export/plugins/megatron_importer.py (1)

515-521: Harden softmax_offset gating to avoid AttributeError/KeyError

Good tighten on the None check. However, hasattr(attention.core_attention, ...) still dereferences attention.core_attention unconditionally and will raise if attention lacks that attribute. Also, indexing self.rules["softmax_offset"] can KeyError when the rule isn’t present for some arch. Guard both.

Apply this diff:

-                        if (
-                            hasattr(attention.core_attention, "softmax_offset")
-                            and attention.core_attention.softmax_offset is not None
-                        ):
-                            self.rules["softmax_offset"](
-                                attention.core_attention.softmax_offset, layer_id
-                            )
+                        core_attn = getattr(attention, "core_attention", None)
+                        softmax_offset = getattr(core_attn, "softmax_offset", None)
+                        softmax_rule = self.rules.get("softmax_offset")
+                        if softmax_offset is not None and callable(softmax_rule):
+                            softmax_rule(softmax_offset, layer_id)

Please confirm all non-MLA attention variants hit here always have core_attention (or this guard will be needed). Also confirm the rule exists for arches where softmax_offset is expected (no unexpected skips).

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

modelopt/torch/speculative/plugins/megatron_eagle.py (1)

1194-1208: Validate offline mode inputs and forbid return_eagle_inputs.

The offline forward path correctly validates required inputs but there's a potential issue with the validation logic.

This matches the existing comment about validating required inputs and forbidding return_eagle_inputs in offline mode.

modelopt/torch/speculative/plugins/megatron_eagle.py (1)

767-767: Clarify the comment about offline mode.

The comment mentions that "layer ids are not used in offline eagle" but this could be more precise. EAGLE-3 still uses layer IDs internally for determining the FC input size multiplier, even in offline mode.

Apply this diff to improve the comment clarity:

-        # layer ids are not used in offline eagle, but we need to set this to have correct fc_input_size_multiplier
+        # layer ids are not actively extracted in offline eagle, but we still need to set this to have correct fc_input_size_multiplier

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

modelopt/torch/speculative/plugins/megatron_eagle.py (2)

1197-1208: ****

Offline forward path: validate required inputs and forbid return_eagle_inputs.

• If eagle_offline=True, aux_hidden_states and hidden_states must be provided; currently missing keys lead to NoneType errors later.
• return_eagle_inputs should be rejected in offline mode to avoid ambiguous semantics.

---

1224-1226: PP early return should be outside the online/offline branching.

The pipeline parallel early return (if not self.post_process) is only in the online path. This prevents offline mode from working correctly with pipeline parallelism.

Move the PP check outside of the conditional to handle both paths:

             return_eagle_inputs=return_eagle_inputs,
         )
 
-        # Typically, this is only the case when PP > 1.
-        if not self.post_process:
-            return hidden_states
+    # Typically, this is only the case when PP > 1.
+    if not self.post_process:
+        return hidden_states

modelopt/torch/speculative/plugins/megatron_eagle.py (2)

1344-1344: Inconsistent indexing for eagle_logits extraction in offline mode.

The code uses logits_sbh.shape[0] for eagle_logits slicing which differs between online and offline modes. Using negative indexing (-labels.shape[1]:) would be more consistent and less error-prone.

Apply consistent negative indexing for all eagle_logits extractions:

-        eagle_logits_1 = eagle_logits_2x[-labels.shape[1] :, :, :]
+        eagle_logits_1 = eagle_logits_2x[-labels.shape[1]:, :, :]
 
-        eagle_logits_2 = eagle_logits_3x[-labels.shape[1] :, :, :]
+        eagle_logits_2 = eagle_logits_3x[-labels.shape[1]:, :, :]
 
-        eagle_logits_3 = eagle_logits_4x[-labels.shape[1] :, :, :]
+        eagle_logits_3 = eagle_logits_4x[-labels.shape[1]:, :, :]

Also applies to: 1386-1386, 1428-1428

---

1282-1292: Replace zero-padding with an ignore-index (or warn) when labels are 1 token short

labels are padded with 0 when labels.shape[1] == input_ids.shape[1] - 1 (modelopt/torch/speculative/plugins/megatron_eagle.py:1282-1292), which will make the last token contribute to loss if logit distillation is disabled. eagle_self_logit_distillation exists in the codebase (passed through config and asserted in the plugin), but the padding branch still runs in some cases — avoid training corruption by masking instead of zero-padding.

• Recommended fix (preferred): pad with the loss ignore index (e.g. -100) or the project’s IGNORE_TOKEN_ID:
  right_token_pad = torch.full((labels.shape[0], 1), -100, dtype=labels.dtype, device=labels.device)
  labels = torch.cat((labels, right_token_pad), dim=-1)
• Alternate (minimum): log a warning when padding happens and logit distillation is disabled:
  if not self.eagle_self_logit_distillation:
  import warnings; warnings.warn("Label padding in offline mode may introduce training errors when logit distillation is disabled.")
• Location: modelopt/torch/speculative/plugins/megatron_eagle.py lines ~1282-1292 (padding) and ~736-761 (eagle_self_logit_distillation use/assert).

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro